from genutils.simple_ga import Population
from genutils.io import from_config, from_config_string
from bEandI_classes_python_only import *
from genutils.smoothing import smooth
from genutils.signal_derivative import signal_derivative_1D
import pylab

class ThisPop(Population):
    def set_scaling_functions(self, begin_index, end_index):
        self.data_begin_index = begin_index
        self.data_end_index   = end_index
        self.scaling_functions = []
        for i in range(len(self.c.data)):
            signal = self.c.data[i].average[begin_index: end_index]
            # smooth the signal
            smoothed_signal = smooth(signal, 20) # use a window width of 20 steps
            # take the first derivative of the signal
            ss_derivative = signal_derivative_1D(smoothed_signal, self.c.data_dt)
            # smooth the derivative and take another
            smoothed_ssd = smooth(signal, 20)
            sssd_derivative = signal_derivative_1D(smoothed_ssd, self.c.data_dt)
            # rectify and then smooth once more.
            self.scaling_functions.append( smooth(abs(sssd_derivative), 50)  )
        
    
    def score_model(self,gene,printme=False):
        values = []
        for allyl in gene.allyls:
            values.append(allyl.value)
        self.c.set_synapses_from_list(values)
        self.c.exercise()
        #residuals = self.c.calculate_residuals()
        difference_arrays = self.c.calculate_residuals(self.data_begin_index, self.data_end_index)
        # multiply the difference arrays by the scaling functions
        differences = []
        for sfn,diffa in zip(self.scaling_functions,difference_arrays):
            diffa *= sfn
            differences.append( sum(diffa) )
        return sum(differences)

    def gene_plausible(self, gene):
        # the onset of I2 should be later than I1
        # return True # always true for now.
        if gene.allyls[0].value > gene.allyls[4].value:
            return False
        return True
        
# read in parameters
cfile = "ONS_ga.parameters"
model_config_filename    = from_config_string('model_config_filename', cfile)
model_parameter_filename = from_config_string('model_parameter_filename', cfile)
num_generations          = int(from_config('num_generations',cfile))
population_size          = int(from_config('population_size',cfile))
initial_population_size  = int(from_config('initial_population_size',cfile))
mutate_rate              = from_config('mutate_rate',cfile)
random_mutate_rate       = from_config('random_mutate_rate',cfile)
cross_rate               = from_config('cross_rate',cfile)

# initalize the genetic algorithm population
thispop = ThisPop(population_size,initial_population_size,cfile)
thispop.mutate_rate = mutate_rate
thispop.cross_rate  = cross_rate

# set up the model
thispop.c = Cell(config_filename=model_config_filename)
# more setting up of the model
thispop.c.set_synapses(model_parameter_filename)
# set up the scaling functions for scoring model performance
thispop.set_scaling_functions(40, 2000) # arguements are the begining and end of relavent data indices.

# run through the generations and print out some info each generation.
for i in range(num_generations):
    thispop.evolve()

    """
    # show the results of all this hard work.
    # ... plot the best result of the final generation
    values = []
    for allyl in thispop.generation[-1][0].allyls:
        values.append(allyl.value)

    thispop.c.set_synapses_from_list(values)
    thispop.c.exercise()
    thispop.c.pylab_plot()
    """

    avg_score = 0
    for gene in thispop.generation[-1]:
        avg_score += gene.score
    avg_score /= len(thispop.generation[-1])

    # print out some info about this generation.
    print  str(i) + ')  Top = ' + str(thispop.generation[-1][0].score) +\
           " Average = " + str(avg_score)
    diversity = thispop.genetic_diversity( thispop.generation[-1] )
    total_diversity = 1
    for d in diversity:
        total_diversity *= d
    print  "    Diversity = " + str( diversity )
    print  "    Total diversity = " + str(total_diversity) 

# archive the genetic population
import cPickle
import random
filename = "ONS_population-" + str( int(random.random()*1000) ) + ".cPickle"
cPickle.dump(thispop.generation, open(filename,'w'), True)
print "Archived population to " + filename

# show the results of all this hard work.
# ... plot the best result of the final generation
values = []
for allyl in thispop.generation[-1][0].allyls:
    values.append(allyl.value)

thispop.c.set_synapses_from_list(values)
thispop.c.exercise()
print thispop
thispop.c.pylab_plot()
import pylab
pylab.show()

    

